from PIL import Image
import random
import os
# 定义 ANSI 转义码
green_background = "\033[48;5;113m"
red_background = "\033[48;5;203m"
yellow_background = "\033[48;5;226m"
blue_background = "\033[48;5;159m"
pink_background = '\033[48;2;255;182;193m'
black_font = "\033[30m"
reset = "\033[0m"

###定义的一些具有背景色的print
def greenbk_print(str,end=None):
    if type(str) != type("123"):
        exit(-1)
    print(green_background+black_font+str+reset,end=end)
def redbk_print(str,end=None):
    if type(str) != type("123"):
        exit(-1)
    print(red_background+black_font+str+reset,end=end)
def yellowbk_print(str,end=None):
    if type(str) != type("123"):
        exit(-1)
    print(yellow_background+black_font+str+reset,end=end)
def bluebk_print(str,end=None):
    if type(str) != type("123"):
        exit(-1)
    print(blue_background+black_font+str+reset,end=end)
def pinkbk_print(str,end=None):
    if type(str) != type("123"):
        exit(-1)
    print(pink_background+black_font+str+reset,end=end)
def string2binstream(string):
    if type(string)!=type('123'):
        redbk_print('[-]Invalid Input!')
        exit(-1)
    string=str(len(string))+':'+string
    binstream=''
    for c in string:
        binstream+=bin(ord(c))[2:].zfill(8)
    return binstream
def binstream2string(binary_str):
    # 检查输入字符串是否只包含 '1' 或 '0'
    string=''
    if not all(bit in '01' for bit in binary_str):
        redbk_print("输入字符串应只包含 '1' 或 '0'.")
        exit(-1)
    for i in range(0,len(binary_str),8):
        string+=chr(int(binary_str[i:i+8],2))
    return string
### 处理RGBA图像的隐写
def RGBA_info(image,binInfo,output_path='string_output.png'):
    '''
    binInfo形如：'010001100011011110101'
    隐写方法为：
    FOR CLOUMN:
        FOR ROW:
            FOR R->G->B->A:
                FOR MASKBITS
                    REWRITE BIT
    '''
    if image.mode!='RGBA':
        redbk_print('[-] Not [RGBA] image-1')
        exit(-1)
    
    ### 创建一个等大小的空白画布，以便进行隐写信息添加
    steg_pic=Image.new('RGBA',image.size,(0,0,0,0))

    ### 记录选择的RGBA隐写位
    bitsMask=[]

    ### 选择RGBA隐写的位
    bluebk_print('[~]Now you are going to choose bits you wanna use to hide info in channels:')
    bluebk_print('[~]You have to input numbers from [0,7] splited by comma like: 0,3,7')
    for channel in ['r','g','b','a']:
        yellowbk_print(f'[~]Channel {channel}:')
        try:
            bluebk_print('[~]Input:',end='')
            bits=input().strip(',')
            if not len(bits):
                bitsMask.append(set())
                continue
            # print(bits)
            bits=bits.split(',')
            for bit in bits:
                if int(bit)<0 or int(bit)>=8:
                    raise ValueError
            bitsMask.append(set(bits))
        except ValueError:
            redbk_print('[-]Invalid Input!')
            exit(-1)
    
    ### 开始隐写
    index=0
    cloumn,row=image.size
    for x in range(cloumn):
        for y in range(row):
            pixel=list(image.getpixel((x,y)))
            for channel in range(4):
                if index>=len(binInfo):
                    break
                ### 取出需要修改的掩码
                for bit in bitsMask[channel]:
                    if int(binInfo[index]):
                        pixel[channel]=pixel[channel]|(1<<int(bit))
                    else:
                        pixel[channel]=pixel[channel]&(255-(1<<int(bit)))
                    index+=1
                    if index>=len(binInfo):
                        break
            steg_pic.putpixel((x,y),tuple(pixel))
    steg_pic.save(output_path)
    greenbk_print(f'[+]Success hide info to {output_path}')
### 处理RGBA图像的隐写
def RGBA_reinfo(image,output_path='output.png',steg_pic=False):
    '''
    binInfo形如：'010001100011011110101'
    隐写方法为：
    FOR CLOUMN:
        FOR ROW:
            FOR R->G->B->A:
                FOR MASKBITS
                    REWRITE BIT
    '''
    if image.mode!='RGBA':
        redbk_print('[-] Not [RGBA] image-2')
        exit(-1)
    ### 记录选择的RGBA隐写位
    bitsMask=[]
    ### 选择RGBA隐写的位
    bluebk_print('[~]Now you are going to choose bits you wanna use to hide info in channels:')
    bluebk_print('[~]You have to input numbers from [0,7] splited by comma like: 0,3,7')
    for channel in ['r','g','b','a']:
        yellowbk_print(f'[~]Channel {channel}:')
        try:
            bluebk_print('[~]Input:',end='')
            bits=input().strip(',')
            if not len(bits):
                bitsMask.append(set())
                continue
            # print(bits)
            bits=bits.split(',')
            for bit in bits:
                if int(bit)<0 or int(bit)>=8:
                    raise ValueError
            bitsMask.append(set(bits))
        except ValueError:
            redbk_print('[-]Invalid Input!')
            exit(-1)
        
    binInfo=''
    ### 开始隐写信息回复
    cloumn,row=image.size
    for x in range(cloumn):
        for y in range(row):
            pixel=list(image.getpixel((x,y)))
            for channel in range(4):
                ### 取出需要掩码对应位的信息
                for bit in bitsMask[channel]:
                    binInfo+=str(int((pixel[channel]&1<<int(bit))!=0))
    string=binstream2string(binInfo)
    colon_index=string.index(':')
    length=int(string[:colon_index])
    string2=string[colon_index+1:]
    string2=string2[:length]
    greenbk_print(f'[+]The hidden info: {string2}')
    
def image2binstream(image):
    if image.mode!='RGB':
        redbk_print('[-] Not [RGB] image')
        exit(-1)
    binstream=''
    # head=str(image.size[0])+':'+str(image.size[1])+'+'
    # for c in head:
    #     binstream+=bin(ord(c))[2:].zfill(8)
    for x in range(image.size[0]):
        for y in range(image.size[1]):
            pixel=list(image.getpixel((x,y)))
            for channel in range(3):
                for bit in range(8):
                    binstream+=str(int((pixel[channel]&1<<bit)!=0))
    # print(binstream[:100])
    return binstream,image.size[0],image.size[1]

def RGBA_pic(image,secret_image,save_path='output2.png'):
    try:
        image=image.convert('RGBA')
    except:
        pass
    if image.mode!='RGBA':
        redbk_print('[-] Not [RGBA] image-3')
        exit(-1)
    if secret_image.mode!='RGB':
        try:
            secret_image=secret_image.convert('RGBA')
        except:
            pass
        try:
            secret_image=secret_image.convert('RGB')
        except:
            redbk_print('[-] Not [RGB] image')
            exit(-1)
    binstream,size0,size1=image2binstream(secret_image)
    len_binstream=len(binstream)
    ### 记录选择的RGBA隐写位
    bitsMask=[]

    ### 选择RGBA隐写的位
    bluebk_print('[~]Now you are going to choose bits you wanna use to hide info in channels:')
    bluebk_print('[~]You have to input numbers from [0,7] splited by comma like: 0,3,7')
    for channel in ['r','g','b','a']:
        yellowbk_print(f'[~]Channel {channel}:')
        try:
            bluebk_print('[~]Input:',end='')
            bits=input().strip(',')
            if not len(bits):
                bitsMask.append(set())
                continue
            # print(bits)
            bits=bits.split(',')
            for bit in bits:
                if int(bit)<0 or int(bit)>=8:
                    raise ValueError
            bitsMask.append(set(bits))
        except ValueError:
            redbk_print('[-]Invalid Input!')
            exit(-1)
    len_bits=sum([len(_) for _ in bitsMask])
    magic=8
    while magic/8*len_binstream>=(len_bits*image.size[0]*image.size[1]):
        magic-=1
    if magic==0:
        redbk_print('[-]No space for hiding!')
        exit(-1)
    # print(magic)
    binInfo=''.join([binstream[i+8-magic:i+8] for i in range(0,len(binstream),8)])
    head=''
    for c in str(size0)+':'+str(size1)+'+':
        head+=bin(ord(c))[2:].zfill(8)
    # print('head: '+head)
    binInfo=head+binInfo

    ### 创建一个等大小的空白画布，以便进行隐写信息添加
    steg_pic=Image.new('RGBA',image.size,(0,0,0,0))
    ### 开始隐写
    index=0
    cloumn,row=image.size
    for x in range(cloumn):
        for y in range(row):
            pixel=list(image.getpixel((x,y)))
            for channel in range(4):
                if index>=len(binInfo):
                    break
                ### 取出需要修改的掩码
                for bit in bitsMask[channel]:
                    if int(binInfo[index]):
                        pixel[channel]=pixel[channel]|(1<<int(bit))
                    else:
                        pixel[channel]=pixel[channel]&(255-(1<<int(bit)))
                    index+=1
                    if index>=len(binInfo):
                        break
            steg_pic.putpixel((x,y),tuple(pixel))
    steg_pic.save(save_path)
    greenbk_print(f'[+]Success hide img to {save_path}')
    greenbk_print('[+]The MAGIC:',end='');yellowbk_print(f'{magic}')

def RGBA_repic(image,magic):
    if image.mode!='RGBA':
        redbk_print('[-] Not [RGBA] image-4')
        exit(-1)

    ### 记录选择的RGBA隐写位
    bitsMask=[]

    ### 选择RGBA隐写的位
    bluebk_print('[~]Now you are going to choose bits you wanna use to hide info in channels:')
    bluebk_print('[~]You have to input numbers from [0,7] splited by comma like: 0,3,7')
    for channel in ['r','g','b','a']:
        yellowbk_print(f'[~]Channel {channel}:')
        try:
            bluebk_print('[~]Input:',end='')
            bits=input().strip(',')
            if not len(bits):
                bitsMask.append(set())
                continue
            # print(bits)
            bits=bits.split(',')
            for bit in bits:
                if int(bit)<0 or int(bit)>=8:
                    raise ValueError
            bitsMask.append(set(bits))
        except ValueError:
            redbk_print('[-]Invalid Input!')
            exit(-1)

    binInfo=''
    ### 开始隐写信息回复
    cloumn,row=image.size
    for x in range(cloumn):
        for y in range(row):
            pixel=list(image.getpixel((x,y)))
            for channel in range(4):
                ### 取出需要掩码对应位的信息
                for bit in bitsMask[channel]:
                    binInfo+=str(int((pixel[channel]&1<<int(bit))!=0))

    # 检查输入字符串是否只包含 '1' 或 '0'
    string=''
    if not all(bit in '01' for bit in binInfo):
        redbk_print("输入字符串应只包含 '1' 或 '0'.")
        exit(-1)
    # print(binInfo[:80])
    for i in range(0,len(binInfo),8):
        c=chr(int(binInfo[i:i+8],2))
        if c=='+':
            break
        string+=chr(int(binInfo[i:i+8],2))
    # return string
    # input(string)
    content=binInfo[i+8:]
    size0,size1=string.split(':');size0=int(size0);size1=int(size1)
    content_length=size0*size1*3*magic
    
    ### 隐写图像回复
    pixel_index=0
    steg_pic=Image.new('RGB',(size0,size1),(0,0,0))
    for x in range(size0):
        for y in range(size1):
            if pixel_index<content_length:
                pixel=content[pixel_index:pixel_index+3*magic];pixel_index+=3*magic
                r=int(''.join(reversed(pixel[:magic].zfill(8))),2)
                g=int(''.join(reversed(pixel[magic+1:2*magic].zfill(8))),2)
                b=int(''.join(reversed(pixel[2*magic:].zfill(8))),2)
                steg_pic.putpixel((x,y),(r,g,b))
            else:
                redbk_print('[-]Error!')
    steg_pic.show()
    greenbk_print(f'[+]Successfully re-image')

def en_main(testflag=True):
    bluebk_print('================================================')
    bluebk_print('[~]Choose the image to hide info:',end='')
    image_path=input()
    if testflag:
        image=Image.open('123.jpg')
    else:
        if not os.path.exists(image_path):
            redbk_print('[-]Error! File not exists!')
            exit(-1)
        image=Image.open(image_path)
    greenbk_print('[+]Successfully open target image')
    bluebk_print('------------------------------------------------')
    if image.mode=='1':
        yellowbk_print(f'[!]Image Mode:{image.mode} [二值图像]')
        redbk_print('[-]二值图像无法使用LSB低位隐写')
        exit(-1)
    elif image.mode=='L':
        greenbk_print(f'[+]Image Mode:{image.mode} [灰度图像]')
        yellowbk_print('[!]对灰度值图像进行LSB隐写可能会效果不佳')
        greenbk_print('[+]正在将图像转换为RGBA模式...')
        image=image.convert('RGBA')
        greenbk_print('[+]转换成功')
        greenbk_print('[+]Image Mode:RGBA [真彩色加透明度]')
    elif image.mode=='P':
        greenbk_print(f'[+]Image Mode:{image.mode} [8位调色板图像]')
        greenbk_print('[+]正在将图像转换为RGBA模式...')
        image=image.convert('RGBA')
        greenbk_print('[+]转换成功')
        greenbk_print('[+]Image Mode:RGBA [真彩色加透明度]')
    elif image.mode=='RGB':
        greenbk_print(f'[+]Image Mode:{image.mode} [真彩色图像]')
        greenbk_print('[+]正在将图像转换为RGBA模式...')
        image=image.convert('RGBA')
        greenbk_print('[+]转换成功')
        greenbk_print('[+]Image Mode:RGBA [真彩色加透明度]')
    elif image.mode=='RGBA':
        greenbk_print(f'[+]Image Mode:{image.mode} [真彩色加透明度')
    else:
        redbk_print('[-]Unknown or Unsupported Mode!')
        exit(-1)
    bluebk_print('------------------------------------------------')
    bluebk_print('[~]Choose the secret-info-mode (0:image/1:string): ',end='')
    secret_switch=input()
    while(secret_switch!='0' and secret_switch!='1'):
        redbk_print('[-]False input! Again')
        bluebk_print('[~]Choose the secret-info-mode (0:image/1:string): ',end='')
        secret_switch=input()
    if secret_switch=='0':
        greenbk_print('[+]You choose to hide secret image!')
        bluebk_print('[+]Enter the path to secret image:',end='')
        if testflag:
            input()
            secret_image_path='cat.png'
        else:
            secret_image_path=input()
            if not os.path.exists(secret_image_path):
                redbk_print('[-]The path you entered is invalid!')
                exit(-1)
        try:
            image2=Image.open(secret_image_path)
            greenbk_print('[+]Successfully open secret image')
        except:
            redbk_print('[-]Could not open secret image')
            input()
            exit(-1)
        bluebk_print('[~]Input the name of output file:',end='')
        if testflag:
            input()
            save_path='string_output.png'
        else:
            save_path=input()
        RGBA_pic(image,image2,save_path=save_path)
    else:
        greenbk_print('[+]You choose to hide secret string!')
        bluebk_print('[~]Enter the string to be hidden:',end='')
        string=input()
        bluebk_print('[~]Input the name of output file:',end='')
        if testflag:
            input()
            secret_string_path='string_output.png'
        else:
            secret_string_path=input()
        binstream=string2binstream(string)
        RGBA_info(image,binstream,output_path=secret_string_path)

def de_main(testflag=True):
    bluebk_print('------------------------------------------------')
    bluebk_print('[~]Choose the image with secret info:',end='')
    image_path=input()
    if testflag:
        image=Image.open('string_output.png')
    else:
        if not os.path.exists(image_path):
            redbk_print('[-]Error! File not exists!')
            exit(-1)
        image=Image.open(image_path)
    greenbk_print('[+]Successfully open target image')
    if image.mode!='RGBA':
        yellowbk_print('[!]Image Mode must be RGBA')
        redbk_print(f'[-]Image Mode:{image.mode}')
        exit(-1)
    bluebk_print('------------------------------------------------')
    bluebk_print('[~]Choose the secret-info-mode (0:image/1:string): ',end='')
    secret_switch=input()
    while(secret_switch!='0' and secret_switch!='1'):
        redbk_print('[-]False input! Again')
        bluebk_print('[~]Choose the secret-info-mode (0:image/1:string): ',end='')
        secret_switch=input()
    if secret_switch=='0':
        if testflag:
            image=Image.open('output2.png')
        greenbk_print('[+]You choose to decrypt secret image!')
        bluebk_print('[~]Please input the magic number: ',end='')
        try:
            magic=int(input())
        except:
            redbk_print('[-]Invalid input!')
            input()
            exit(-1)
        RGBA_repic(image,magic)
    else:
        greenbk_print('[+]You choose to decrypt secret string!')
        RGBA_reinfo(image)

class colors:
    PURPLE = '\033[95m'
    CYAN = '\033[96m'
    DARKCYAN = '\033[36m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    END = '\033[0m'
def print_colored(text,end='\n'):
    color = random.choice([colors.PURPLE, colors.CYAN, colors.DARKCYAN, colors.BLUE, colors.GREEN, colors.YELLOW, colors.RED])
    pinkbk_print(color + text + colors.END,end=end)
def welcome():
    Title='''
 __       __                      __                  __         ______   _______            ______    __                         
/  \     /  |                    /  |                /  |       /      \ /       \          /      \  /  |                        
$$  \   /$$ |  ______    ______  $$/   _______       $$ |      /$$$$$$  |$$$$$$$  |        /$$$$$$  |_$$ |_     ______    ______  
$$$  \ /$$$ | /      \  /      \ /  | /       |      $$ |      $$ \__$$/ $$ |__$$ | ______ $$ \__$$// $$   |   /      \  /      \ 
$$$$  /$$$$ | $$$$$$  |/$$$$$$  |$$ |/$$$$$$$/       $$ |      $$      \ $$    $$< /      |$$      \$$$$$$/   /$$$$$$  |/$$$$$$  |
$$ $$ $$/$$ | /    $$ |$$ |  $$ |$$ |$$ |            $$ |       $$$$$$  |$$$$$$$  |$$$$$$/  $$$$$$  | $$ | __ $$    $$ |$$ |  $$ |
$$ |$$$/ $$ |/$$$$$$$ |$$ \__$$ |$$ |$$ \_____       $$ |_____ /  \__$$ |$$ |__$$ |        /  \__$$ | $$ |/  |$$$$$$$$/ $$ \__$$ |
$$ | $/  $$ |$$    $$ |$$    $$ |$$ |$$       |      $$       |$$    $$/ $$    $$/         $$    $$/  $$  $$/ $$       |$$    $$ |
$$/      $$/  $$$$$$$/  $$$$$$$ |$$/  $$$$$$$/       $$$$$$$$/  $$$$$$/  $$$$$$$/           $$$$$$/    $$$$/   $$$$$$$/  $$$$$$$ |
                       /  \__$$ |                                                                                       /  \__$$ |
                       $$    $$/                                                                                        $$    $$/ 
                        $$$$$$/                                                                                          $$$$$$/  
'''
    print_colored(Title)

if __name__ == '__main__':
    welcome()
    print_colored('### Welcome to Magic LSB-Steg tools ###')
    print_colored('### Create by HKcoder 2024.7 ###')
    while(1):
        print_colored('================================================')
        print_colored('[@]Please choose to en-steg or de-steg [en:0/de:1]:',end='')
        choice=input().strip()
        while(choice!='0' and choice!='1'):
            redbk_print('[-]Input invaild')
            print_colored('[@]Please choose to en-steg or de-steg [en:0/de:1]:',end='')
            choice=input()
        if choice=='0':
            en_main(testflag=False)
        else:
            de_main(testflag=False)
